/*****************************************************
 * @file  selector.cpp
 * @brief 模拟和验证select网络选择器功能。
 *****************************************************/

#if defined(PlantOS)
#include "init.h"
#endif
#include <atomic>
#include <iostream>
#include <mutex>
#include <thread>
#include <unordered_map>
#include <vector>

// ============== PlantOS ==============
#if defined(PlantOS)
#include <PlantOS.h>
#include "arpa/inet.h"
#include "sockLib.h"
#include "inetd.h"
#include "hostUtils.h"
#include <unistd.h>
#include <sysTypes.h>
//  ============== Windows ==============
#elif defined(_WIN32) || defined(_WIN64)
#include <winsock2.h>
#include <ws2tcpip.h>
#pragma comment(lib, "ws2_32.lib")
#define intYPE SOCKET
#define close(sock) closesocket(sock)
#define FD_SET(n, p) ((p)->fds_bits[(n) / BITSPERWORD] |= ((fd_mask) 1 << ((n) % BITSPERWORD)))
#define FD_CLR(n, p) ((p)->fds_bits[(n) / BITSPERWORD] &= ~((fd_mask) 1 << ((n) % BITSPERWORD)))
#define FD_ISSET(n, p) ((p)->fds_bits[(n) / BITSPERWORD] & ((fd_mask) 1 << ((n) % BITSPERWORD)))
#define FD_ZERO(p) bzero((char *) (p), sizeof(*(p)))
#else
// ==============  Others  ==============
#include <arpa/inet.h>
#include <fcntl.h>
#include <unistd.h>
#endif

#define BUFFER_SIZE 1024
#define SELECT_TIMEOUT 1

using namespace std;

class Context;
using listen_handler = void (*)(int fd, const char *buffer, ssize_t len);

struct Listener {
	int            sockfd;
	uint16_t       port;
	listen_handler handler;
	Listener(int s = -1, uint16_t p = 0) : sockfd(s), port(p), handler(nullptr) {}
	~Listener() {
		if (sockfd != -1) {
			close(sockfd);
			cout << "[Listenor] close socket: " << sockfd << endl;
		}
	}
};

class Context {
private:
	unordered_map<int, unique_ptr<Listener>> sock2listener;
	unordered_map<uint16_t, Listener *>      port2listener;
	mutex                                    _spps_lock;

public:
	void addListener(unique_ptr<Listener> listener) {
		lock_guard<mutex> lock(_spps_lock);
		auto              ptr = listener.get();
		sock2listener.emplace(ptr->sockfd, move(listener));
		port2listener.emplace(ptr->port, ptr);
		cout << "[Context] Add listentor: " << ptr->sockfd << " port: " << ptr->port << endl;
	}

	Listener *getListenerByFD(int fd) {
		lock_guard<mutex> lock(_spps_lock);
		auto              it = sock2listener.find(fd);
		return it != sock2listener.end() ? it->second.get() : nullptr;
	}
	Listener *getListenerByPort(uint16_t port) {
		lock_guard<mutex> lock(_spps_lock);
		auto              it = port2listener.find(port);
		return it != port2listener.end() ? it->second : nullptr;
	}

	void removeListener(uint16_t port) {
		lock_guard<mutex> lock(_spps_lock);
		auto              port_it = port2listener.find(port);
		if (port_it != port2listener.end()) {
			auto *ptr = port_it->second;
			cout << "[Context] remove listentor: " << ptr->sockfd << " port: " << ptr->port << endl;
			sock2listener.erase(ptr->sockfd);
			port2listener.erase(port_it);
		}
	}

	vector<int> getAllSockets() {
		lock_guard<mutex> lock(_spps_lock);
		vector<int>       sockets;
		for (const auto &pair: sock2listener) {
			sockets.push_back(pair.first);
		}
		return sockets;
	}
};

class Selector {
private:
	atomic<bool> running{false};
	thread       worker_thread;
	Context     &context;

	void worker() {
		while (running) {
			auto sockets = context.getAllSockets();
			if (sockets.empty()) {
				this_thread::sleep_for(chrono::seconds(1));
				continue;
			}

			fd_set readfds;
			FD_ZERO(&readfds);
			int max_fd = 0;
			for (auto sock: sockets) {
				FD_SET(sock, &readfds);
				max_fd = max(max_fd, sock);
			}

			struct timeval timeout{.tv_sec = SELECT_TIMEOUT, .tv_usec = 0};
			int            ret = select(max_fd + 1, &readfds, nullptr, nullptr, &timeout);
			if (ret < 0) {
				cout << "[Selector] select error." << endl;
				this_thread::sleep_for(chrono::seconds(3));
				continue;
			}

			for (auto sock: sockets) {
				if (FD_ISSET(sock, &readfds)) {
					char               buffer[BUFFER_SIZE];
					struct sockaddr_in client_addr{};
					socklen_t          addr_len = sizeof(client_addr);
					ssize_t            len      = recvfrom(sock, buffer, BUFFER_SIZE, 0, (sockaddr *) &client_addr, &addr_len);
					if (len > 0) {
						if (auto *listener = context.getListenerByFD(sock)) {
							listener->handler(sock, buffer, len);
						} else if (len < 0 && (errno == EBADF || errno == ENOTSOCK)) {
							if (auto *listener = context.getListenerByFD(sock)) {
								context.removeListener(listener->port);
							}
						}
					}
				}
			}
		}
	}

public:
	Selector(Context &ctx) : context(ctx) {
		running       = true;
		worker_thread = thread(&Selector::worker, this);
	}
	~Selector() {
		running = false;
		if (worker_thread.joinable()) {
			worker_thread.join();
		}
	}
};

unique_ptr<Listener> createListener(uint16_t port = 0) {
	int sockfd = socket(AF_INET, SOCK_DGRAM, IPPROTO_UDP);
	if (sockfd < 1) {
		cerr << "[Listenor] create error." << endl;
		return nullptr;
	}

	// 地址重用
	int reuse = 1;
	if (setsockopt(sockfd, SOL_SOCKET, SO_REUSEADDR, &reuse, sizeof(reuse)) < 0) {
		cerr << "[Listenor] reuse setting failed." << endl;
		return nullptr;
	}


	struct sockaddr_in server_addr{};
	server_addr.sin_family      = AF_INET;
	server_addr.sin_addr.s_addr = htonl(INADDR_ANY);
	server_addr.sin_port        = htons(port);
	if (::bind(sockfd, (sockaddr *) &server_addr, sizeof(server_addr)) == -1) {
		cerr << "[Listenor] bind failed。" << endl;
		close(sockfd);
		return nullptr;
	}

	socklen_t len = sizeof(server_addr);
	getsockname(sockfd, (sockaddr *) &server_addr, &len);
	auto actual_port = ntohs(server_addr.sin_port);
	cout << "[Listenor] create successed - socket: " << sockfd << " port: " << actual_port << endl;

	return unique_ptr<Listener>(new Listener(sockfd, actual_port));
}

void udp_handler(int fd, const char *buffer, ssize_t len) {
	//
	cout << "[UDPHandler]  fd:" << fd << ",  data: " << string(buffer, len) << endl;
}

// g++ -std=c++11 -pthread -o selector.exe selector.cpp
int main() {
	// 创建Context（包含listener集合）
	Context          context;
	vector<uint16_t> managed_port;
	for (int i = 0; i < 5; ++i) {
		auto listener = createListener();
		if (listener) {
			listener->handler = udp_handler;
			managed_port.push_back(listener->port);
			context.addListener(move(listener));
		}
	}
	std::cout << endl;

	// Selector线程监听listener集合
	Selector selector(context);

	// 模拟客户端
	int count = 1;
	for (int i = 0; i < 3; i++) {
		this_thread::sleep_for(chrono::milliseconds(1));
		for (auto port: managed_port) {
			if (auto *listener = context.getListenerByPort(port)) {
				string             msg = "hello_" + to_string(count);
				struct sockaddr_in target{};
				target.sin_family      = AF_INET;
				target.sin_addr.s_addr = htonl(INADDR_LOOPBACK);
				target.sin_port        = htons(port);
				sendto(listener->sockfd, msg.data(), msg.size(), 0, (sockaddr *) &target, sizeof(target));
				count++;
			}
		}
	}

	// 回车结束程序，并执行析构程序
	getchar();
	cout << "\n [Cleanup] start clean..." << endl;
	for (auto port: managed_port) {
		context.removeListener(port);
	}
	cout << "[Cleanup] clean finished..." << endl;
	return 0;
}
